import numpy as np
import onnxruntime
from PIL import Image


def preprocess_image(image_path):
    """
    将任意分辨率的图像加载、灰度化、缩放到 28x28 并归一化，以符合 MNIST 模型输入要求。
    """
    # 加载图像并转换为灰度
    img = Image.open(image_path).convert("L")
    # 缩放图像到 28x28
    img = img.resize((28, 28))
    # 转换为 numpy 数组并归一化到 [0, 1]
    img_array = np.array(img, dtype=np.float32) / 255.0
    # 将图像调整为 (1, 1, 28, 28) 的形状，以匹配模型的输入
    img_array = img_array.reshape(1, 1, 28, 28)
    return img_array


# 加载 ONNX 模型
onnx_session = onnxruntime.InferenceSession(
    "/home/lurker/文档/playground/python/pytorch/mnist_model.onnx"
)


def predict(image_path):
    """
    使用 ONNX 模型对图像进行预测。
    """
    # 预处理图像
    input_data = preprocess_image(image_path)

    # 准备输入
    input_name = onnx_session.get_inputs()[0].name
    output_name = onnx_session.get_outputs()[0].name
    prediction = onnx_session.run([output_name], {input_name: input_data})

    # 获取预测类别
    predicted_class = np.argmax(prediction[0])
    print(f"预测类别: {predicted_class}")


# 使用任意分辨率图像的路径
image_path = (
    "/home/lurker/文档/playground/python/pytorch/test/test0.png"  # 待推理图片路径
)
predict(image_path)
